import gym
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn.functional as 火炬函数
from torch import nn, optim
import matplotlib

matplotlib.rc("font", family='Microsoft YaHei')

埃普西隆 = 0.9
伽马 = 0.9
学习率 = 0.01
记忆体_容量 = 3000
质量_网络_迭代数 = 100
每批数量 = 32

插曲数 = 400
是否渲染 = True
if 是否渲染:
    游戏环境 = gym.make("MountainCar-v0", render_mode="human")
else:
    游戏环境 = gym.make("MountainCar-v0")
游戏环境 = 游戏环境.unwrapped
状态数 = 游戏环境.observation_space.shape[0]  # 汽车沿x轴的位置，汽车的速度
动作数 = 游戏环境.action_space.n  # 3个动作，（向左，向右，不动）


class 网络(nn.Module):
    def __init__(self):
        super(网络, self).__init__()

        self.全连接层1 = nn.Linear(状态数, 30)
        self.全连接层1.weight.data.normal_(0, 0.1)
        self.全连接层2 = nn.Linear(30, 动作数)
        self.全连接层2.weight.data.normal_(0, 0.1)

    def forward(self, x):
        x = self.全连接层1(x)
        x = 火炬函数.relu(x)
        x = self.全连接层2(x)

        return x


class 深度质量网络:
    def __init__(self):
        self.评估网络, self.目标网络 = 网络(), 网络()
        self.记忆体 = np.zeros((记忆体_容量, 状态数 * 2 + 2))
        #
        self.记忆次数 = 0
        self.学习次数 = 0
        self.优化器 = optim.Adam(self.评估网络.parameters(), 学习率)
        self.损失函数 = nn.MSELoss()

        self.图像, self.轴 = plt.subplots()

    def 选择动作(self, 状态):
        状态 = torch.unsqueeze(torch.FloatTensor(状态), 0)
        # 使其具备探索的能力
        if np.random.randn() <= 埃普西隆:
            动作得分 = self.评估网络.forward(状态)
            动作 = torch.max(动作得分, 1)[1].data.numpy()
            动作 = 动作[0]
        else:
            动作 = np.random.randint(0, 动作数)

        return 动作

    def 存储(self, 状态, 动作, 奖励, 下一个状态):
        if self.记忆次数 % 500 == 0:
            print("经验池收集{}次经验".format(self.记忆次数))
        索引 = self.记忆次数 % 记忆体_容量
        数据 = np.hstack((状态, [动作], [奖励], 下一个状态))
        self.记忆体[索引,] = 数据
        self.记忆次数 += 1

    def 学习(self):
        if self.学习次数 % 质量_网络_迭代数 == 0:
            self.目标网络.load_state_dict(self.评估网络.state_dict())
        self.学习次数 += 1

        样本索引 = np.random.choice(记忆体_容量, 每批数量)
        每批_记忆 = self.记忆体[样本索引, :]
        每批_状态 = torch.FloatTensor(每批_记忆[:, :状态数])

        每批_动作 = torch.LongTensor(每批_记忆[:, 状态数:状态数 + 1].astype(int))
        每批_奖励 = torch.FloatTensor(每批_记忆[:, 状态数 + 1:状态数 + 2])
        每批_下一个_状态 = torch.FloatTensor(每批_记忆[:, -状态数:])

        质量_评估 = self.评估网络(每批_状态)
        质量_评估 = 质量_评估.gather(1, 每批_动作)
        质量_下一个 = self.目标网络(每批_下一个_状态).detach()
        a = 质量_下一个.max(1)
        质量_目标 = 每批_奖励 + 伽马 * 质量_下一个.max(1)[0].view(每批数量, 1)

        损失值 = self.损失函数(质量_评估, 质量_目标)
        self.优化器.zero_grad()
        损失值.backward()
        self.优化器.step()

    def 绘制(self, 轴, x):
        x = np.array(x)
        轴.cla()
        轴.set_xlabel("插曲")
        轴.set_ylabel("完整奖励")
        轴.plot(x[:, 0], x[:, 1], 'b-')  # b：蓝色的。-：实线样式
        plt.pause(0.000000000000001)


def 主要():
    一个网络 = 深度质量网络()
    print("深度质量网络正在收集经验。")
    步数列表 = []
    for 插曲 in range(插曲数):
        状态, _ = 游戏环境.reset()
        步数 = 0
        while True:
            步数 += 1
            if 是否渲染:
                游戏环境.render()
            动作 = 一个网络.选择动作(状态)
            下一个状态, 奖励, 已完毕, _, _ = 游戏环境.step(动作)
            奖励 = 奖励 * 100 if 奖励 > 0 else 奖励 * 5
            步数列表.append((步数, 奖励))

            一个网络.存储(状态, 动作, 奖励, 下一个状态)
            if 一个网络.记忆次数 >= 记忆体_容量:
                一个网络.学习()
                if 已完毕:
                    print("插曲 {}，这个奖励是{}。".format(插曲, round(奖励, 3)))

            if 已完毕:
                一个网络.绘制(一个网络.轴, 步数列表)
                break
            状态 = 下一个状态


if __name__ == '__main__':
    主要()
